import json
import pandas as pd

from datetime import datetime, timedelta
from xunkemgmt.service.market_quotes_service_pb2 import \
    QueryKlinesRequestProto, QueryCurrencyExchangeRatesRequestProto
from coin2.service.strategy.trading_summary_service_pb2 import \
    QueryTradingSummaryIntervalHistoriesRequestProto
from coin2.service.account import \
    account_info_service_pb2 as account_info_pb2
import coin.proto.coin_request_pb2 as cr_pb2
from xunkemgmt_client.client.api_client import XunkemgmtClient
from coin.proto.coin_market_enums_pb2 import KlineInterval
from google.protobuf.json_format import MessageToDict
from coin.base.datetime_util import to_timestamp_int
from coin.util.protobuf.util import protos_to_df


_pta_rename_mapping = {
    'volume_mark': 'fill_qty',
    'turnover_mark': 'fill_pq',
    'turnover_mark_in_usd': 'fill_pq_in_usd',
    'turnover_maker_mark': 'fill_pq_maker',
    'turnover_taker_mark': 'fill_pq_taker',
    'turnover_maker_mark_in_usd': 'fill_pq_maker_in_usd',
    'turnover_taker_mark_in_usd': 'fill_pq_taker_in_usd',
    'turnover_maker_inferred_mark_in_usd': 'fill_pq_maker_inferred_in_usd',
    'turnover_taker_inferred_mark_in_usd': 'fill_pq_taker_inferred_in_usd',
    'turnover_maker_inferred_mark': 'fill_pq_maker_inferred',
    'turnover_taker_inferred_mark': 'fill_pq_taker_inferred',
    'pnl_mark_from_fill': 'pnl_from_fill',
    'pnl_mark_from_fill_in_usd': 'pnl_from_fill_in_usd',
    'pnl_mark_from_balance_inferred': 'pnl_from_balance',
    'pnl_mark_from_balance_inferred_in_usd': 'pnl_from_balance_in_usd'
}


def _gen_kline_dict(kline_data):
  kline_dict = {
    'symbol': kline_data.get('symbol', None),
    'native_symbol': kline_data.get('native_symbol', None),
    'market_type': kline_data.get('market_type', None),
    'exchange': kline_data.get('exchange', None),
    'kline_period': kline_data.get('kline_period', None),
    'klines': kline_data['klines']
  }
  kline_dict = {k: v for k, v in kline_dict.items() if v}
  return json.dumps(kline_dict)


def query_klines(start_date: datetime,
                 end_date: datetime,
                 *,
                 market_types=None,
                 exchanges=None,
                 symbols=None,
                 period='KLINE_INTERVAL_1DAY',
                 as_proto=False):
  assert period in KlineInterval.keys(), period
  response_list = []
  trading_date = start_date
  with XunkemgmtClient() as client:
    while trading_date <= end_date:
      next_date = trading_date + timedelta(days=1)
      kline_request = QueryKlinesRequestProto(
          start_timestamp=to_timestamp_int(trading_date),
          end_timestamp=to_timestamp_int(trading_date + timedelta(days=1)),
          market_types=market_types,
          exchanges=exchanges,
          symbols=symbols,
          kline_period=period,
      )
      trading_date = next_date
      response = client.query_klines(kline_request)
      klines = getattr(response, 'klines', [])
      response_list.extend(klines)
  if as_proto:
    return response_list
  else:
    kline_df = protos_to_df(response_list)
    if len(kline_df) > 0:
      kline_df['trading_date'] = kline_df['trading_date'].apply(
          lambda x: pd.to_datetime(str(x)))
      kline_df['symbol_contract_type'] = kline_df['relative_symbol']
      kline_df['kline_dict'] = kline_df.apply(lambda x: _gen_kline_dict(x), axis=1)
    return kline_df


def query_pta(start_date,
              end_date,
              *,
              business_units,
              strategy_ids=None,
              strategy_names=None,
              strategy_groups=None,
              traders=None,
              account_ids=None,
              market_types=None,
              exchanges=None,
              owners=None,
              symbols=None,
              asset_names=None,
              pta_type=None,
              as_proto=False):
  pta_request = QueryTradingSummaryIntervalHistoriesRequestProto(
      start_timestamp=to_timestamp_int(start_date),
      end_timestamp=to_timestamp_int(end_date + timedelta(days=1)),
      business_units=business_units,
      strategy_ids=strategy_ids,
      strategy_names=strategy_names,
      strategy_groups=strategy_groups,
      traders=traders,
      account_ids=account_ids,
      market_types=market_types,
      exchanges=exchanges,
      owners=owners,
      symbols=symbols,
      asset_names=asset_names,
  )
  with XunkemgmtClient() as client:
    response = client.query_trading_summary_interval_histories(pta_request)
  strat_stats = getattr(response, 'histories', [])
  if as_proto:
    return strat_stats
  else:
    response_list = []
    for stat in strat_stats:
      stat = {
          **MessageToDict(stat.summary_info,
                          preserving_proto_field_name=True),
          **MessageToDict(stat.summary,
                          preserving_proto_field_name=True)
      }
      response_list.append(stat)
    pta_df = pd.DataFrame(response_list)
    if len(pta_df) > 0:
      if 'symbol' not in pta_df.columns:
        pta_df['type'] = 'STRATEGY'
      else:
        pta_df['type'] = pta_df['symbol'].apply(
            lambda x: 'STRATEGY' if pd.isnull(x) else 'SYMBOL')
      if pta_type:
        pta_df = pta_df[pta_df['type'] == pta_type]
      pta_df['trading_date'] = pta_df['trading_date'].apply(
          lambda x: pd.to_datetime(str(x)))
      pta_df.rename(_pta_rename_mapping, axis=1, inplace=True)
  return pta_df


def query_exchange_rates(coinlist, quote, trading_date=None):
  if trading_date:
    trading_date = int(trading_date.strftime('%Y%m%d'))
  symbol_list = [f'{coin.upper()}-{quote.upper()}' for coin in coinlist]
  ticker_request = QueryCurrencyExchangeRatesRequestProto(
      symbols=symbol_list,
      trading_date=trading_date,
  )
  with XunkemgmtClient() as client:
    response = client.query_currency_exchange_rates(ticker_request)
  exchange_rates = getattr(response, 'exchange_rates', [])
  exchange_rates_dict = {}
  for rate in exchange_rates:
    last = rate.last if rate.HasField('last') else None
    exchange_rates_dict[rate.symbol] = last
  exchange_rats = [exchange_rates_dict.get(symbol, None) for symbol in symbol_list]
  return exchange_rats


def query_exchange_tickers(*,
                           query_ts=None,
                           exchange_api_id=None,
                           market_type=None,
                           exchange=None,
                           api_version=None):
  with XunkemgmtClient() as client:
    data = client.query_exchange_tickers(
        cr_pb2.ExchangeApiRequestProto(
            query_ts=query_ts,
            exchange_api_id=exchange_api_id,
            market_type=market_type,
            exchange=exchange,
            api_version=api_version,
    ))
  return data


def query_transfer_histories(start_time,
                             end_time,
                             *,
                             transfer_type=None,
                             transfer_status=None):
  start_timestamp = to_timestamp_int(start_time) \
    if isinstance(start_time, datetime) else start_time
  end_timestamp = to_timestamp_int(end_time) \
    if isinstance(end_time, datetime) else end_time
  with XunkemgmtClient() as client:
    data = client.query_transfer_histories(
        account_info_pb2.QueryTransferHistoriesRequestProto(
            start_timestamp=start_timestamp,
            end_timestamp=end_timestamp,
            transfer_type=transfer_type,
            transfer_status=transfer_status,
    )).histories
  return data
